package cn.edu.hitsz.compiler.asm;

import cn.edu.hitsz.compiler.ir.IRVariable;

import java.util.*;

public class RegAllocator {

    private final List<Reg> availableRegs;
    private final Set<Reg> usedRegs;
    private final BMap<IRVariable, Reg> variableToRegMap;

    public RegAllocator() {
        // 初始化寄存器池，假设 t0 到 t6 是临时寄存器
        this.availableRegs = Arrays.asList(
                new Reg("t0", 0), new Reg("t1", 1), new Reg("t2", 2),
                new Reg("t3", 3), new Reg("t4", 4), new Reg("t5", 5),
                new Reg("t6", 6)
        );
        this.usedRegs = new HashSet<>();
        this.variableToRegMap = new BMap<>();
    }

    // 分配一个寄存器给变量
    public Reg allocateRegForVariable(IRVariable variable) {
        if (variableToRegMap.containsKey(variable)) {
            return variableToRegMap.getByKey(variable);  // 如果已经分配，则直接返回
        }
        for (Reg reg : availableRegs) {
            if (!usedRegs.contains(reg)) {
                usedRegs.add(reg);
                variableToRegMap.put(variable, reg);
                return reg;
            }
        }
        // 如果没有可用寄存器，抛出异常或进行溢出处理
        throw new RuntimeException("No available registers for allocation");
    }


    // 获取已分配给特定变量的寄存器
    public Reg getRegForVariable(IRVariable variable) {
        return variableToRegMap.getByKey(variable);
    }

    // 释放一个寄存器（用于回收未使用的变量寄存器）
    public void freeReg(Reg reg) {
        usedRegs.remove(reg);
        variableToRegMap.removeByValue(reg);
    }
}
